{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env OMP_NUM_THREADS=10\n",
    "%pylab inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import os\n",
    "import sys\n",
    "import cv2\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pylab as plt\n",
    "import tensorflow as tf\n",
    "import inspect\n",
    "from collections import namedtuple\n",
    "import time\n",
    "import glob\n",
    "from imageio import imread, imsave\n",
    "print('USE GPU #{} Host={}'.format(os.environ.get('CUDA_VISIBLE_DEVICES'), os.uname()[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOCAL_PATH = '../'\n",
    "if LOCAL_PATH not in sys.path:\n",
    "    sys.path.append(LOCAL_PATH)\n",
    "from common.tf_layer_utils import *\n",
    "from mydatasets import *\n",
    "from det_tools import *\n",
    "MODEL_PATH = '../models/'\n",
    "if MODEL_PATH not in sys.path:\n",
    "    sys.path.append(MODEL_PATH)\n",
    "COMMON_PATH = '../common/'\n",
    "if COMMON_PATH not in sys.path:\n",
    "    sys.path.append(COMMON_PATH)\n",
    "    \n",
    "from io_utils import read_text\n",
    "from jupyter_utils import display_image_batch # make sure you finish preparation described at jupyter_utils.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Loading and Image warping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Outdoor scenes (YFCC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yfcc_dir = '../release/outdoor_examples/'\n",
    "img_dir = os.path.join(yfcc_dir, 'images')\n",
    "dpt_dir = os.path.join(yfcc_dir, 'depths')\n",
    "seq_names = ['sacre_coeur']\n",
    "batch_size = 8\n",
    "data_raw_size = 362 # data_size * sqrt(2)\n",
    "data_size = 256\n",
    "aug_max_degree = 180\n",
    "aug_max_scale = 1.0\n",
    "depth_thresh = 10.0 # depend on dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "data_loader = SfMDataset(out_size=(data_raw_size, data_raw_size),\n",
    "                         warp_aug_mode='random', flip_pair=True,\n",
    "                         max_degree=aug_max_degree, max_scale=aug_max_scale,\n",
    "                        )\n",
    "dataset = data_loader.get_dataset(dpt_dir, img_dir, seq_names, phase='train',\n",
    "                                  batch_size=batch_size, shuffle=True, seed=1234)\n",
    "next_batch = dataset.make_one_shot_iterator().get_next()\n",
    "next_batch = euclidean_augmentation(next_batch, (data_size, data_size), rot_aug=True, scale_aug=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "photos1, photos2, depths1, depths2, valid_masks1, valid_masks2, c2Tc1s, c1Tc2s, c1Tws, c2Tws, Ks1, Ks2, thetas1, thetas2, inv_thetas1, inv_thetas2, theta_params = next_batch\n",
    "batch, height, width, _ = tf.unstack(tf.shape(photos1))\n",
    "\n",
    "photos1w, visible_masks1, xy_maps1to2 = \\\n",
    "    inverse_warp_view_2_to_1(photos2, depths2, depths1, c2Tc1s, \n",
    "                            K1=Ks1, K2=Ks2, \n",
    "                            inv_thetas1=inv_thetas1, thetas2=thetas2,\n",
    "                            depth_thresh=depth_thresh)\n",
    "photos2w, visible_masks2, xy_maps2to1 = \\\n",
    "    inverse_warp_view_2_to_1(photos1, depths1, depths2, c1Tc2s, \n",
    "                            K1=Ks2, K2=Ks1,\n",
    "                            inv_thetas1=inv_thetas2, thetas2=thetas1,\n",
    "                            depth_thresh=depth_thresh)\n",
    "visible_masks1 = visible_masks1 * valid_masks1 # take 'and'\n",
    "visible_masks2 = visible_masks2 * valid_masks2\n",
    "\n",
    "diff1 = tf.abs(visible_masks1*(photos1-photos1w))\n",
    "diff2 = tf.abs(visible_masks2*(photos2-photos2w))\n",
    "\n",
    "canvas1 = tf.concat([photos1, photos1w, diff1], axis=2)\n",
    "canvas2 = tf.concat([photos2, photos2w, diff2], axis=2)\n",
    "canvas = tf.concat([canvas1, canvas2], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfconfig = tf.ConfigProto()\n",
    "tfconfig.gpu_options.allow_growth = True \n",
    "sess = tf.Session(config=tfconfig)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fetch_dict = {\n",
    "    'canvas': canvas,\n",
    "}\n",
    "\n",
    "outs = sess.run(fetch_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_image_batch(outs['canvas'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View keypoints from dumped npz\n",
    "You need to run run_lfnet.py first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_name = 'london_bridge-00000.jpg'\n",
    "img_path = '../samples/'+img_name\n",
    "kp_path = '../dump_feats/{}.npz'.format(img_name)\n",
    "image = imread(img_path).astype(np.float32) / 255\n",
    "kp_info = np.load(kp_path)\n",
    "kpts = kp_info['kpts']\n",
    "feats = kp_info['descs']\n",
    "height, width = kp_info['size']\n",
    "image = cv2.resize(image, (width, height))\n",
    "print('Load {}-kpts, feat={}-D imsize={}x{}'.format(kpts.shape[0], feats.shape[1], width, height))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(image)\n",
    "plt.plot(\n",
    "    kpts[:,0], kpts[:,1],\n",
    "    alpha=1,\n",
    "    linestyle='none',\n",
    "    linewidth=0,\n",
    "    aa=False,\n",
    "    marker='.',\n",
    "    markersize=2,\n",
    "    color=[0.9, 0.1, 0.1],\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Indoor scenes (ScanNet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scannet_dir = '/cvlabdata1/home/ono/datasets/scannet/dataset/all/valid' # set your own path\n",
    "render_text = '/cvlabdata1/home/ono/datasets/scannet/dataset/valid_nips.txt'\n",
    "render_paths = [x for x in read_text(render_text)]\n",
    "\n",
    "batch_size = 8\n",
    "pair_offset = 30\n",
    "random_offset = False\n",
    "data_raw_size = 362 # data_size * sqrt(2)\n",
    "data_size = 256\n",
    "aug_max_degree = 180\n",
    "aug_max_scale = 1.0\n",
    "depth_thresh = 1.0 # depend on dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "data_loader = SE3PairwiseDataset(batch_size, offset_val=pair_offset, random_offset=random_offset,\n",
    "                out_width=data_raw_size, out_height=data_raw_size, \n",
    "                crop_center=True,\n",
    "                max_degree=aug_max_degree, max_scale=aug_max_scale, warp_aug_mode='random')\n",
    "data_loader.set_files(scannet_dir, render_paths, max_img_num=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = data_loader.get_dataset(shuffle=True)\n",
    "next_batch = dataset.make_one_shot_iterator().get_next()\n",
    "next_batch = euclidean_augmentation(next_batch, (data_size, data_size), rot_aug=True, scale_aug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "photos1, photos2, depths1, depths2, valid_masks1, valid_masks2, c2Tc1s, c1Tc2s, c1Tws, c2Tws, Ks1, Ks2, thetas1, thetas2, inv_thetas1, inv_thetas2, theta_params = next_batch\n",
    "batch, height, width, _ = tf.unstack(tf.shape(photos1))\n",
    "\n",
    "photos1w, visible_masks1, xy_maps1to2 = \\\n",
    "    inverse_warp_view_2_to_1(photos2, depths2, depths1, c2Tc1s, \n",
    "                            K1=Ks1, K2=Ks2, \n",
    "                            inv_thetas1=inv_thetas1, thetas2=thetas2,\n",
    "                            depth_thresh=depth_thresh)\n",
    "photos2w, visible_masks2, xy_maps2to1 = \\\n",
    "    inverse_warp_view_2_to_1(photos1, depths1, depths2, c1Tc2s, \n",
    "                            K1=Ks2, K2=Ks1,\n",
    "                            inv_thetas1=inv_thetas2, thetas2=thetas1,\n",
    "                            depth_thresh=depth_thresh)\n",
    "visible_masks1 = visible_masks1 * valid_masks1 # take 'and'\n",
    "visible_masks2 = visible_masks2 * valid_masks2\n",
    "\n",
    "diff1 = tf.abs(visible_masks1*(photos1-photos1w))\n",
    "diff2 = tf.abs(visible_masks2*(photos2-photos2w))\n",
    "\n",
    "canvas1 = tf.concat([photos1, photos1w, diff1], axis=2)\n",
    "canvas2 = tf.concat([photos2, photos2w, diff2], axis=2)\n",
    "canvas = tf.concat([canvas1, canvas2], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfconfig = tf.ConfigProto()\n",
    "tfconfig.gpu_options.allow_growth = True \n",
    "sess = tf.Session(config=tfconfig)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fetch_dict = {\n",
    "    'canvas': canvas,\n",
    "}\n",
    "\n",
    "outs = sess.run(fetch_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outs['canvas'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(outs['canvas'][0,...,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
